from neo4j import GraphDatabase

# Neo4j数据库连接配置
URI = "bolt://localhost:7687"
AUTH = ("neo4j", "12345678")

def query_herbal_medicine(driver, medicine_name):
    """查询中药及其关联内容（匹配实际标签和关系）"""
    query = """
    MATCH (m:中药 {name: $name})-[r]-(related)
    RETURN 
        m.name as medicine_name,
        m.formual as formula,
        m.effect as effect,
        type(r) as relation_type,
        labels(related) as related_labels,
        related.name as related_name,
        related.type as related_type
    """
    with driver.session() as session:
        result = session.run(query, name=medicine_name)
        return [dict(record) for record in result]

def format_result(record):
    """格式化查询结果"""
    relation_map = {
        "治疗": "可治疗疾病",
        "作用于": "适用症状",
        "包含": "组成药材"  # 新增材料关系显示
    }
    return {
        "relation_type": relation_map.get(record["relation_type"], record["relation_type"]),
        "related_type": record["related_type"],
        "related_name": record["related_name"]
    }

def run_medicine_query(uri, auth):
    """中药知识图谱查询交互"""
    driver = GraphDatabase.driver(uri, auth=auth)
    try:
        print("中药知识图谱查询系统 (输入 'exit' 退出)")
        while True:
            name = input("\n请输入中药/方剂名称:\n> ").strip()

            if name.lower() in ("exit", "quit"):
                break

            if not name:
                print("错误：名称不能为空")
                continue

            try:
                records = query_herbal_medicine(driver, name)
                if not records:
                    print(f"未找到与 '{name}' 相关的内容")
                    continue

                # 提取基础信息
                first_record = records[0]
                print(f"\n【{name}】详细信息:")
                print(f"- 组成: {first_record.get('formula', '无记录')}")
                print(f"- 功效: {first_record.get('effect', '无记录')}")

                # 打印关联关系
                print(f"\n关联内容 ({len(records)} 条):")
                for i, record in enumerate(map(format_result, records), 1):
                    print(f"{i}. [{record['related_type']}] {record['related_name']} ({record['relation_type']})")

            except Exception as e:
                print(f"查询错误: {type(e).__name__}: {e}")

    finally:
        driver.close()
        print("\n数据库连接已关闭")

if __name__ == "__main__":
    run_medicine_query(URI, AUTH)